package main

import (
	"context"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"sync"
	"time"
)

// Client 结构体封装 WebSocket 连接，保证写操作的线程安全
type Client struct {
	Conn *websocket.Conn
	mu   sync.Mutex
}

var (
	clients   sync.Map            // 存储所有在线的客户端
	broadcast = make(chan []byte) // 消息广播通道
)

// WebSocket 升级器
var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

func main() {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	http.HandleFunc("/chat", handleChat)
	go handleMessages(ctx)

	log.Println("Server started on localhost:8808")
	if err := http.ListenAndServe(":8808", nil); err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}

// 处理 WebSocket 连接
func handleChat(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("Upgrade error:", err)
		return
	}
	client := &Client{Conn: conn}
	clientID := conn.RemoteAddr().String()
	clients.Store(clientID, client)
	log.Printf("New connection: %s", clientID)

	defer func() {
		clients.Delete(clientID)
		conn.Close()
		log.Printf("Connection closed: %s", clientID)
	}()

	// 设置 Pong 处理，避免超时断连
	conn.SetReadDeadline(time.Now().Add(60 * time.Second))
	conn.SetPongHandler(func(string) error {
		conn.SetReadDeadline(time.Now().Add(60 * time.Second))
		return nil
	})

	// 发送心跳 Ping
	go func() {
		ticker := time.NewTicker(30 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				client.mu.Lock()
				err := client.Conn.WriteMessage(websocket.PingMessage, nil)
				client.mu.Unlock()
				if err != nil {
					log.Println("Ping error:", err)
					conn.Close()
					clients.Delete(clientID)
					return
				}
			}
		}
	}()

	for {
		_, message, err := conn.ReadMessage()
		if err != nil {
			break
		}
		log.Printf("Received message from %s: %s", clientID, string(message))
		broadcast <- message
	}
}

// 处理消息广播
func handleMessages(ctx context.Context) {
	for {
		select {
		case message := <-broadcast:
			distributeMessage(message)
		case <-ctx.Done():
			return
		}
	}
}

// 分发消息给所有在线客户端
func distributeMessage(message []byte) {
	clients.Range(func(key, value interface{}) bool {
		client := value.(*Client)
		client.mu.Lock()
		defer client.mu.Unlock()

		if err := client.Conn.WriteMessage(websocket.TextMessage, message); err != nil {
			log.Println("WriteMessage error:", err)
			client.Conn.Close()
			clients.Delete(key)
		}
		return true
	})
}
